package tower.ai.base.chat.config;

import org.springframework.ai.chroma.ChromaApi;
import org.springframework.ai.ollama.OllamaEmbeddingModel;
import org.springframework.ai.vectorstore.ChromaVectorStore;
import org.springframework.ai.vectorstore.VectorStore;
import org.springframework.beans.factory.annotation.Value;
import org.springframework.boot.autoconfigure.AutoConfigureAfter;
import org.springframework.boot.autoconfigure.condition.ConditionalOnProperty;
import org.springframework.context.annotation.Bean;
import org.springframework.context.annotation.Configuration;
import org.springframework.context.annotation.Primary;
import org.springframework.http.client.SimpleClientHttpRequestFactory;
import org.springframework.web.client.RestClient;

/**
 * model: 查看OllamaOptions类, 里面的模型表示支持的模型
 */
@Configuration
@AutoConfigureAfter(OllamaConfig.class)
@ConditionalOnProperty(prefix = "spring.chroma", name = "enabled", havingValue = "true", matchIfMissing = false)
public class ChromaConfig {
    @Value("${spring.chroma.url}")
    private String chromaUrl;
    @Bean
    public RestClient.Builder builder() {
        return RestClient.builder().requestFactory(new SimpleClientHttpRequestFactory());
    }

    //    目前调用chromaDB接口不通
    @Bean("ollamaVectorStore")
    @Primary
    public VectorStore ollamaVectorStore(OllamaEmbeddingModel ollamaEmbeddingModel, ChromaApi chromaApi) {
        return new ChromaVectorStore(ollamaEmbeddingModel, chromaApi, "chenkw-ollama-embedding", true);
    }

    @Bean
    public ChromaApi chromaApi(RestClient.Builder restClientBuilder) {
        ChromaApi chromaApi = new ChromaApi(chromaUrl, restClientBuilder);
        return chromaApi;
    }
}
